Skip to content

PYTHON-5672 Refactor connection checkout to reduce layers#2893

Open
blink1073 wants to merge 20 commits into
mongodb:mainfrom
blink1073:PYTHON-5672
Open

PYTHON-5672 Refactor connection checkout to reduce layers#2893
blink1073 wants to merge 20 commits into
mongodb:mainfrom
blink1073:PYTHON-5672

Conversation

@blink1073

@blink1073 blink1073 commented Jun 25, 2026

Copy link
Copy Markdown
Member

PYTHON-5672

Changes in this PR

Flatten the connection checkout hot path by replacing @asynccontextmanager generator-based context managers with class-based ones, and eliminate _MongoClientErrorHandler by absorbing it into the new _ClientCheckout class. Reduces per-checkout overhead by ~3 generator frames.

Test Plan

  • just typing passes (mypy + pyright)
  • just lint-manual passes (all pre-commit hooks)
  • Existing CMAP and unified format spec tests cover the checkout/checkin contract and SDAM error handling ordering

Checklist

Checklist for Author

  • Did you update the changelog (if necessary)?
  • Is there test coverage?
  • Is any followup work tracked in a JIRA ticket? If so, add link(s). PYTHON-5898 - Reduce server selection and pool lock overhead on the checkout hot path

Checklist for Reviewer

  • Does the title of the PR reference a JIRA Ticket?
  • Do you fully understand the implementation? (Would you be comfortable explaining how this code works to someone else?)
  • Is all relevant documentation (README or docstring) updated?

Replace the three @asynccontextmanager layers on the connection checkout
hot path with class-based async context managers, and eliminate
_MongoClientErrorHandler by absorbing it into _ClientCheckout.

- _PoolCheckout replaces Pool.checkout() generator CM
- _ClientCheckout replaces _checkout() generator CM and absorbs all of
  _MongoClientErrorHandler (contribute_socket, handle, SDAM error logic)
- _ClientReadCheckout extends _ClientCheckout to apply single-topology
  read preference adjustment (formerly _conn_from_server())
- active_contexts.add() consolidated into _get_conn(), avoiding a
  separate lock acquisition on the hot path; deduped so new connections
  (already tracked by connect()) are not double-added
- Connection leak fixed: self._conn assigned before event publishing so
  checkin runs if a CMAP listener raises in __aenter__
- _ClientCheckout.for_existing_conn() classmethod handles the
  _run_operation() getMore path that needs SDAM handling around an
  already-checked-out connection
…onnectionRetryable

Eliminates per-instance __dict__ allocation and defers the _deprioritized_servers
list creation until it is actually needed (only on sharded retry paths).
@codecov-commenter

codecov-commenter commented Jun 25, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 89.83051% with 24 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
pymongo/asynchronous/mongo_client.py 85.71% 9 Missing and 1 partial ⚠️
pymongo/synchronous/mongo_client.py 85.71% 9 Missing and 1 partial ⚠️
pymongo/asynchronous/pool.py 95.83% 1 Missing and 1 partial ⚠️
pymongo/synchronous/pool.py 95.83% 1 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

blink1073 added 14 commits June 25, 2026 13:03
…nter__

When _PoolCheckout.__aenter__() raises (e.g. connect() fails), Python does
not call _ClientCheckout.__aexit__(), so handle() was never invoked and the
topology never learned about the failure. This caused test_5_check_out_fails_
connection_error to fail (missing PoolClearedEvent) and broke failover tests.

Fix: wrap both the pool checkout call and the post-checkout setup in
try/except BaseException blocks inside _ClientCheckout.__aenter__(), calling
handle() and (for post-checkout failures) checking the connection back in
before re-raising.
Constructing _PoolCheckout directly bypassed pool.checkout() overrides
used in MockPool, causing mock_down_hosts checks to be skipped and
network-error tests to fail.
…rver.checkout()

server.checkout() was a thin wrapper around pool.checkout() that nothing
called after _ClientCheckout was changed to call server.pool.checkout()
directly.

Add a test that mocks publish_connection_checked_out to raise and verifies
the connection is returned to the pool rather than leaked.
…connection

Covers the except BaseException block in _ClientCheckout.__aenter__() that
fires when post-checkout work (session pinning, ConfigurationError check)
raises — verifying the connection is returned to the pool and not leaked.
…thon

patch.object cannot shadow a method on an instance when __slots__ is
defined — the attribute is read-only. Use a subclass override instead.
Resolve import conflict in server.py: upstream removed datetime import
(run_operation timing moved into run_cursor_command), our branch removed
AbstractAsyncContextManager (server.checkout() was deleted). Drop both
now-unused imports.
- Remove _checkout_started_time from _PoolCheckout.__slots__: the value
  was written then immediately aliased to a local variable and never read
  from the slot again, wasting an allocation on every checkout.

- Guard pool_checkout.__aexit__ with try/finally in both _ClientCheckout
  __aexit__ and the __aenter__ setup-failure path: if handle() raises,
  the connection is now guaranteed to be checked in instead of leaked.

- Pass exc.__traceback__ instead of None when calling
  pool_checkout.__aexit__ from the __aenter__ setup-failure except block.
- Reset conn.pinned_txn/pinned_cursor before calling pool_checkout.__aexit__
  in the setup-failure except block so _PoolCheckout.__aexit__ calls checkin()
  instead of silently skipping it when session._pin() ran before the error.
- Replace _ClientReadCheckout._effective_read_pref slot with _read_preference
  and a local variable in __aenter__ to avoid mutating instance state.
- Fix Pool.checkout() docstring to say "with-statement" so synchro.py copies
  the correct text to the sync mirror.
@blink1073
blink1073 marked this pull request as ready for review July 22, 2026 11:52
@blink1073
blink1073 requested a review from a team as a code owner July 22, 2026 11:52
@blink1073
blink1073 requested review from aclark4life and Copilot July 22, 2026 11:52

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Refactors PyMongo’s connection checkout hot path to reduce overhead by replacing generator-based @contextmanager / @asynccontextmanager usage with class-based context managers, and by folding _MongoClientErrorHandler responsibilities into a new _ClientCheckout context manager for both sync and async clients.

Changes:

  • Replace pool checkout generator context managers with _PoolCheckout classes (sync + async) and adjust client/server checkout plumbing accordingly.
  • Introduce _ClientCheckout / _ClientReadCheckout to unify pool checkout, SDAM error handling, and session pinning (sync + async).
  • Add regression tests to ensure connections are not leaked when CMAP event publishing fails or when post-checkout setup raises.

Reviewed changes

Copilot reviewed 14 out of 14 changed files in this pull request and generated 7 comments.

Show a summary per file
File Description
test/test_read_preferences.py Updates read-preference test client override to match new checkout return types.
test/test_pooling.py Adds regression test ensuring checkout doesn’t leak connections if CMAP listener publish fails.
test/test_client.py Adds regression test ensuring client checkout setup failures return connections to the pool.
test/pymongo_mocks.py Updates MockPool checkout override to match new class-based pool checkout behavior.
test/asynchronous/test_read_preferences.py Async equivalent of read-preference override update.
test/asynchronous/test_pooling.py Async equivalent of CMAP publish-failure no-leak regression test.
test/asynchronous/test_client.py Async equivalent of client checkout setup-failure no-leak regression test.
test/asynchronous/pymongo_mocks.py Async equivalent of MockPool checkout override update.
pymongo/synchronous/server.py Removes now-unneeded server-level checkout wrapper in favor of direct pool checkout usage.
pymongo/synchronous/pool.py Implements class-based _PoolCheckout and adjusts pool checkout internals for reuse/new-connection paths.
pymongo/synchronous/mongo_client.py Introduces _ClientCheckout / _ClientReadCheckout and re-routes operation execution through them.
pymongo/asynchronous/server.py Removes now-unneeded async server-level checkout wrapper.
pymongo/asynchronous/pool.py Async implementation of class-based _PoolCheckout.
pymongo/asynchronous/mongo_client.py Async implementation of _ClientCheckout / _ClientReadCheckout and updated operation execution flow.

Comment thread pymongo/synchronous/pool.py Outdated
Comment thread pymongo/synchronous/pool.py Outdated
Comment thread pymongo/synchronous/mongo_client.py Outdated
Comment thread pymongo/synchronous/mongo_client.py Outdated
Comment thread pymongo/asynchronous/pool.py Outdated
Comment thread pymongo/synchronous/mongo_client.py Outdated
Comment thread pymongo/asynchronous/mongo_client.py Outdated
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants